#!/usr/bin/env python3

"""VQ algorithm for image compression

Reference: https://www.sciencedirect.com/topics/engineering/vector-quantization
"""

import numpy as np
from sklearn.mixture import *
from sklearn.cluster import *
from sklearn.base import TransformerMixin

from PIL import Image

class ImageEncoder(TransformerMixin):

    def make_codebook(self):  
        if hasattr(self, 'cluster_centers_'):
            self.codebook_ = self.cluster_centers_
        elif hasattr(self, 'means_'):
            self.codebook_ = self.means_

    def _decode(self, codes):
        """Decode `codes` to the real data with the codebook
        
        Args:
            codes (array-like): array of codes
        
        Returns:
            array: decoding result
        """
        return np.array([self.codebook_[k] for k in codes])

    def decode(self, codes, shape=None):
        # decode to image, a wrapper of _decode
        Y = self._decode(codes)
        if self.block_size == (1, 1):
            Y = Y.reshape((im.size[1], im.size[0], self.n_channels))
            return Image.fromarray(Y.astype('uint8'))
        else:
            return data2image(Y, block_size=self.block_size, shape=shape or self.shape)

    def query(self, k):
        w, h = self.block_size
        return self.codebook_[k].reshape((h, w, self.n_channels))

    def encode(self, im, shape=None):
        X, shape = image2data(im, block_size=self.block_size, shape=shape or self.shape)
        return self.predict(X), shape

    def reconstruct(self, im=None):
        if im is None:
            code, shape = self.labels_, self.shape
        else:
            code, shape = self.encode(im)
        return self.decode(code, shape)

    def stat(self, im, shape=None):
        Y, _ = self.encode(im, shape=shape)
        import collections
        f = collections.Counter(Y)
        N = len(Y)
        p = {k: (v / N) for k, v in f.items()}
        return f, p

    def make_codebook(self):
        if hasattr(self, 'cluster_centers_'):
            self.codebook_ = self.cluster_centers_
        else:
            raise 'You have to define `codebook_`'


class BlockImageEncoder(ImageEncoder):
    """Base class for image encoding by pixel blocks
    """

    def __init__(self, block_size=(1, 1), n_channels=3, **kwargs):
        super().__init__(**kwargs)
        self.block_size = block_size
        self.n_channels = n_channels

    def _fit(self, X, *args, **kwargs):
        super().fit(X, *args, **kwargs)
        self.make_codebook()
        return self


    def fit(self, im, shape=None, *args, **kwargs):
        # fit the image and biuld the codebook
        X, shape = image2data(im, block_size=self.block_size, shape=shape)
        self.shape = shape
        return self._fit(X, *args, **kwargs)


class VQTransformer(BlockImageEncoder, KMeans):
    """VQ algorithm for image compression

    Example:
        model = KMeans(n_clusters=2)
        im = Image.open('cloth.jpg')
        t = VQTransformer(model=model)
        t.fit(im)
        im = t.transform(im)
        im.save('cloth.kmeans2.jpg')
    """

    def __init__(self, block_size=(1, 1), n_channels=3, **kwargs):
        super().__init__(**kwargs)
        self.block_size = block_size
        self.n_channels = n_channels


from semi_kmeans import *

class SemiVQTransformer(BlockImageEncoder, SemiKMeans):
    """Semi VQ algorithm for image compression
    """

    def fit(self, im, mask, shape=None, *args, **kwargs):
        X, shape = image2data(im, block_size=self.block_size, shape=shape)
        self.shape = shape
        _mask = mask.ravel()
        y0 = _mask[_mask != -1]
        X0 = X[_mask!=-1]
        X1 = X[_mask==-1]
        super().fit(X0, y0, X1, *args, **kwargs)
        self.make_codebook()
        n0 = y0.shape[0]
        self.labels_[_mask==-1], self.labels_[_mask!=-1] = self.labels_[n0:], y0
        return self


def make_transformer(cls):
    class Transformer(BlockImageEncoder, cls):
        pass
    return Transformer


from sklearn.cluster import MeanShift
class MeanShiftTransformer(BlockImageEncoder, MeanShift):
    """image compression by meanshift
    """

    def fit(self, im, shape=None, with_pos=False, *args, **kwargs):
        # fit the image and biuld the codebook
        X, shape = image2data(im, block_size=self.block_size, shape=shape, with_pos=with_pos)
        self.shape = shape
        self.with_pos = with_pos
        return self._fit(X, *args, **kwargs)

    def encode(self, im, shape=None):
        X, shape = image2data(im, block_size=self.block_size, shape=shape or self.shape, with_pos=self.with_pos)
        return self.predict(X), shape

    def _decode(self, codes):
        """Decode `codes` with the codebook
        
        Args:
            codes (array-like): array of codes
        
        Returns:
            array: decoding result
        """
        if self.with_pos:
            return np.array([self.codebook_[k][2:] for k in codes])
        else:
            return np.array([self.codebook_[k] for k in codes])


def image2data(im, block_size=(1, 1), shape=None, with_pos=False):
    """Transform an image into 2D array, r*c blocks of w*h array
    where block_size = (w, h)
    each array corresponds to a w*h patch/block of the image
    
    Arguments:
        im {Image} -- an image
    
    Keyword Arguments:
        block_size {tuple} -- a rectangle domain of an image (default: (1, 1))
                        block_size[0] * block_size[1] is the number of features
        shape {tuple} -- size of block matrix (default: {None})
                        calculated based on the size of the image by default
                        shape[0]*shape[1] is the number of samples
    
    Returns:
        tuple of data and shape=(r, c)
    """

    w, h = block_size
    if shape is None:
        W, H = im.size
        c, cc = divmod(W, w)
        r, rr = divmod(H, h)
        im = im.resize((w*c, h*r))
        shape = r, c
    else:
        r, c = shape
        im = im.resize((0, 0, w*c, h*r))
    I = np.asarray(im, dtype=np.float64)
    # blocks <- concatenate the columns of the block matrix
    if with_pos:
        blocks = [np.hstack([[i,j], I[i*h:(i+1)*h, j*w:(j+1)*w].ravel()]) for i in range(r) for j in range(c)]
    else:
        blocks = [I[i*h:(i+1)*h, j*w:(j+1)*w].ravel() for i in range(r) for j in range(c)]
    return np.array(blocks), shape


# def image2data(im, block_size=(1,1), shape=None):
#     blocks, shape = image2blockmatrix(im, block_size=(1,1), shape=None)
#     data = np.array([blocks[i,j].ravel() for j in range(c) for i in range(r)])
#     return data, shape


def data2image(X, block_size, shape):
    # inverse of image2data
    r, c = shape
    w, h = block_size
    X = X.reshape((r, c, w*h*3))
    X = np.block([[[X[i, j].reshape((h, w, 3))] for j in range(c)] for i in range(r)])
    return Image.fromarray(X.astype('uint8', copy=False))


if __name__ == '__main__':

    im = Image.open('lenna.jpg')

    block_size = 4, 4
    n_clusters = 3
    # vq = VQTransformer(block_size=block_size, n_clusters=n_clusters)
    # ms = MeanShiftTransformer(block_size=block_size, max_iter=100)
    # ms.fit(im)

    def demo(model, name='Encoder'):
        im0 = model.reconstruct()
        n_clusters = len(model.codebook_)

        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(1,2)
        ax[0].imshow(im)
        ax[0].set_title('Original Image')
        ax[1].imshow(im0)
        ax[1].set_title(name)
        ax[0].axis('off')
        ax[1].axis('off')
        # fig.suptitle(f'{name} Demo (block-size: {block_size[0]} X {block_size[1]}, #clusters: {n_clusters})')
        plt.show()
    
    def demo_meanshift():
        ms = MeanShiftTransformer(block_size=block_size, max_iter=100)
        ms.fit(im)
        im0 = ms.reconstruct()
        ms.fit(im, with_pos=True)
        im1 = ms.reconstruct()

        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(1,3)
        ax[0].imshow(im)
        ax[0].set_title('Original Image')
        ax[1].imshow(im0)
        ax[1].set_title('MeanShift')
        ax[2].imshow(im1)
        ax[2].set_title('MeanShift (with position)')
        for a in ax:
            a.axis('off')
        # fig.suptitle(f'Demo of MeanShift (block-size: {block_size[0]} X {block_size[1]})')
        plt.show()

    demo_meanshift()

    def _mask(im, block_size=(1,1), shape=None):
        """Transform an image to 2d array
        where block_size = (w, h)
        
        Arguments:
            im {Image} -- an image
        
        Keyword Arguments:
            block_size {tuple} -- a rectangle domain of an image (default: (1, 1))
            shape {tuple} -- size of small domain (default: {None})
                            calculated based on the size of the image by default
        
        Returns:
            tuple of a block matrix and its shape=(r, c)
        """
        w, h = block_size
        if shape is None:
            W, H = im.size
            c, cc = divmod(W, w)
            r, rr = divmod(H, h)
            shape = r, c
            im = im.resize((w*c, h*r))
        else:
            r, c = shape
            im = im.resize((0, 0, w*c, h*r))
        I = np.asarray(im, dtype=np.float64)
        shape = (r, c)
        # blocks <- concatenate the columns of the block matrix
        def _t(a):
            def _f(x):
                if np.all(x == (255,0,0)): return 0
                elif np.all(x == (0,255,0)): return 1
                elif np.all(x == (0,0,255)): return 2
                else: return -1
            return np.max(np.apply_along_axis(_f, 2, a))
        blocks = [[_t(I[i*h:(i+1)*h, j*w:(j+1)*w]) for j in range(c)] for i in range(r)]
        return np.array(blocks)

    def semivq():
        svq = SemiVQTransformer(block_size=block_size, n_clusters=n_clusters)
        masked = Image.open('lenna-masked.jpg')
        mask = _mask(masked, block_size=block_size)
        svq.fit(im, mask=mask)
        im1 = svq.reconstruct()

        import matplotlib.pyplot as plt
        fig, ax = plt.subplots(2, 2)

        ax[0,0].imshow(im)
        ax[0,0].set_title('Original Image')
        ax[0,1].imshow(masked)
        ax[0,1].set_title('Masked Image')
        ax[1,0].imshow(im0)
        ax[1,0].set_title('VQ')
        ax[1,1].imshow(im1)
        ax[1,1].set_title('Semisupervised VQ')

        ax[0,0].axis('off')
        ax[0,1].axis('off')
        ax[1,0].axis('off')
        ax[1,1].axis('off')

        plt.show()
